import os

from langchain_community.chat_models import ChatOpenAI
from langchain_core.tools import tool, Tool

from langchain.agents import initialize_agent, AgentType
from langchain_core.prompts import ChatPromptTemplate
from langchain_ollama import OllamaLLM
from langchain.agents import AgentExecutor, create_tool_calling_agent
from langchain_ollama.chat_models import ChatOllama


os.environ['SERPAPI_API_KEY'] = '47afe0f70fefbe12e10919ee52248ac01d28652b763975bc84347a774805f3b6'

# @tool
# # 示例工具：股票信息查询工具
def get_stock_info(query: str) -> str:
    """用于查询股票的实时价格和基本信息。"""
    return f"股票 {query} 当前价格约为 150 美元。"

stock_tool = Tool(
    name="股票查询工具",
    func=get_stock_info,
    description="用于查询股票的实时价格和基本信息。"
)



# 初始化 Ollama 模型
ollama_llm = ChatOllama(model="MFDoom/deepseek-r1-tool-calling:8b")
#ollama_llm = ChatOllama(model="llama2")
def agent_search_history():
    tools = [stock_tool]
    prompt = ChatPromptTemplate.from_messages(
        [
            ("system", "你是一个金融分析专家，回答要求精准且简明。"),
            ("placeholder", "{chat_history}"),
            ("human", "{input}"),
            ("placeholder", "{agent_scratchpad}"),
        ]
    )

    #ollama_llm.bind_tools([get_stock_info])
    agent = create_tool_calling_agent(ollama_llm, tools, prompt)
    agent_executor = AgentExecutor(agent=agent, tools=tools)

    result = agent_executor.invoke({"input": '请查询苹果公司的股票信息，并解释当前的市场趋势。'})


    # agent_executor = initialize_agent(
    #     tools=tools,
    #     llm=ollama_llm,
    #     agent=AgentType.ZERO_SHOT_REACT_DESCRIPTION,
    #     prompt = prompt,
    #     verbose=True
    # )
    # result = agent_executor.invoke(
    #     {
    #         "input": "请查询苹果公司的股票信息，并解释当前的市场趋势。",
    #         # 如果有历史对话，可以通过 chat_history 传入
    #     }
    # )
    print(result)

if __name__ == '__main__':
    agent_search_history()
